import faiss
import numpy as np

dim = 768
corpus_size = 1000
# np.random.seed(111)

corpus = np.random.random((corpus_size, dim)).astype('float32')

if __name__ == '__main__':
    # 3. Create Index on CPU
    # Faiss 通过直接初始化提供了大量的索引选择
    # 首先构建一个 flat index （在 CPU 上）
    index = faiss.IndexFlatIP(dim)
    # 除了基本的索引类，我们还可以使用 index_factory 函数来生成复合 Faiss 索引
    index = faiss.index_factory(dim, "Flat", faiss.METRIC_L2)

    # 4. Build GPU Index and Search
    """
    All the GPU indexes are built with StandardGpuResources object. It contains all the needed resources for each GPU in use. By default it will allocate 18% of the total VRAM as a temporary scratch space.

    The GpuClonerOptions and GpuMultipleClonerOptions objects are optional when creating index from cpu to gpu. They are used to adjust the way the GPUs stores the objects.
    """

    # use a single GPU
    rs = faiss.StandardGpuResources()
    co = faiss.GpuClonerOptions()
    # then make it to gpu index
    index_gpu = faiss.index_cpu_to_gpu(provider=rs, device=0, index=index, options=co)
    
    index_gpu.add(corpus)
    D, I = index_gpu.search(corpus, 4) # total: 11.6 ms

    # All Available GPUs
    """
    If your system contains multiple GPUs, Faiss provides the option to deploy al available GPUs. You can control their usages through GpuMultipleClonerOptions, e.g. whether to shard or replicate the index acrross GPUs.
    
    如果您的系统包含多个 GPU，Faiss 会提供部署所有可用 GPU 的选项。您可以通过 GpuMultipleClonerOptions 控制它们的用法，例如，是分片还是复制索引 acrross GPU
    """
    # cloner options for multiple GPUs
    co = faiss.GpuMultipleClonerOptions()

    index_gpu = faiss.index_cpu_to_all_gpus(index=index, co=co)
    index_gpu.add(corpus)
    D, I = index_gpu.search(corpus, 4) # total: 56.6 ms

    # Multiple GPUs
    # There's also option that use multiple GPUs but not all:
    ngpu = 4
    resources = [faiss.StandardGpuResources() for _ in range(ngpu)]

    # Create vectors for the GpuResources and divices, then pass them to the index_cpu_to_gpu_multiple() function.
    # 为 GpuResources 和 divice 创建向量，然后将它们传递给 index_cpu_to_gpu_multiple() 函数。
    vres = faiss.GpuResourcesVector()
    vdev = faiss.Int32Vector()
    for i, res in zip(range(ngpu), resources):
        vdev.push_back(i)
        vres.push_back(res)
    index_gpu = faiss.index_cpu_to_gpu_multiple(vres, vdev, index)
    index_gpu.add(corpus)
    D, I = index_gpu.search(corpus, 4) # total: 16.9 ms

    # 5. Results
    # All the three approaches should lead to identical result. Now let's do a quick sanity check:
    # 所有三种方法都应该导致相同的结果。现在让我们做一个快速的健全性检查：
    # The nearest neighbor of each vector in the corpus is itself 
    # 语料库中每个向量的最近邻域是其自身
    assert np.all(corpus[:] == corpus[I[:, 0]])

    print(D[:3])
